from tqdm import tqdm

from yan_new import CvFo
import torch
import pandas as pd
import numpy as np
from eval_multi_data import eval_egg


def train_egg():
    voc = pd.read_pickle("voc_data.pandas_pickle")
    net = CvFo(len(voc), 32, 8, 1, "egg")
    loss_func = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(net.parameters(), lr=0.0001)
    data_set = pd.read_pickle("train_data.pandas_pickle")[:1000]
    batch_size = 12
    epochs = 10
    bar = tqdm(len(epochs * list(range(0, len(data_set), batch_size))))
    for epoch in range(epochs):
        np.random.shuffle(data_set)
        for i in range(0, len(data_set), batch_size):
            j = i + batch_size
            one_data = data_set[i:j]
            two_data = torch.Tensor(one_data).int()
            out = net(two_data[:, :-1])
            loss = loss_func(out.reshape([-1, out.shape[-1]]), two_data[:, 1:].reshape([-1]).long())
            bar.set_description("loss___{:.5f}".format(loss.item()))

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        net.save_egg()


def train_division():
    voc = pd.read_pickle("voc_data.pandas_pickle")
    net = CvFo(len(voc), 32, 8, 1, "division")

    loss_func = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(net.parameters(), lr=0.0001)
    data_sets = eval_egg()
    for division_id, data_set in enumerate(data_sets):
        batch_size = 12
        epochs = 10
        bar = tqdm(len(epochs * list(range(0, len(data_set), batch_size))))
        for epoch in range(epochs):
            np.random.shuffle(data_set)
            for i in range(0, len(data_set), batch_size):
                j = i + batch_size
                one_data = data_set[i:j]
                two_data = torch.Tensor(one_data).int()
                out = net(two_data[:, :-1])
                loss = loss_func(out.reshape([-1, out.shape[-1]]), two_data[:, 1:].reshape([-1]).long())
                bar.set_description("loss___{:.5f}".format(loss.item()))
                bar.update(1)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
        net.save_division(division_id)


def train_dif():
    voc = pd.read_pickle("voc_data.pandas_pickle")
    net = CvFo(len(voc), 32, 8, 5, "differentiation")

    loss_func = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(net.parameters(), lr=0.0001)
    data_set = pd.read_pickle("train_data.pandas_pickle")[:1000]
    batch_size = 12
    epochs = 10
    bar = tqdm(len(epochs * list(range(0, len(data_set), batch_size))))
    for epoch in range(epochs):
        np.random.shuffle(data_set)
        for i in range(0, len(data_set), batch_size):
            j = i + batch_size
            one_data = data_set[i:j]
            two_data = torch.Tensor(one_data).int()
            out = net(two_data[:, :-1])
            loss = loss_func(out.reshape([-1, out.shape[-1]]), two_data[:, 1:].reshape([-1]).long())
            bar.set_description("loss___{:.5f}".format(loss.item()))
            bar.update(1)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        net.save_differentiation()


if __name__ == '__main__':
    train_egg()
    train_division()
    train_dif()
